package eu.gegenleitner.esspacc.xmlSignature;

import java.io.ByteArrayInputStream;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.StringReader;
import java.security.PublicKey;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathExpression;
import javax.xml.xpath.XPathFactory;

import org.apache.xml.security.keys.KeyInfo;
import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.utils.Base64;
import org.apache.xml.security.utils.Constants;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.xml.sax.InputSource;

public class XMLSignatureVerifier {
	
	/**
	 * Loads XML from String
	 * @param xml
	 * @return
	 * @throws Exception
	 */
	private Document loadXMLFromString(String xml, boolean namespaceAware) throws Exception {
	    DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
	    if(namespaceAware) {
	    	factory.setNamespaceAware(true);
	    }
	    DocumentBuilder builder = factory.newDocumentBuilder();
	    xml = xml.replaceFirst("<sl:CreateXMLSignatureResponse xmlns:sl=\"http://www.buergerkarte.at/namespaces/securitylayer/1.2#\">", "");
	    xml = xml.replaceFirst("</sl:CreateXMLSignatureResponse>", "");
	    InputSource is = new InputSource(new StringReader(xml));
	    return builder.parse(is);
	}
	
	/**
	 * Validate the given XMLSignature
	 * @param xmlSignature
	 * @return true if valid, else false
	 */
	public boolean isXMLSignatureValid(String xmlSignature) {
		com.sun.org.apache.xml.internal.security.Init.init();
		org.apache.xml.security.Init.init();
		boolean schemaValidate = true;
        final String signatureSchemaFile = "xmldsig-core-schema.xsd";

        javax.xml.parsers.DocumentBuilderFactory dbf =
            javax.xml.parsers.DocumentBuilderFactory.newInstance();

        if (schemaValidate) {
            dbf.setAttribute("http://apache.org/xml/features/validation/schema",
                             Boolean.TRUE);
            dbf.setAttribute("http://apache.org/xml/features/dom/defer-node-expansion",
                             Boolean.TRUE);
            dbf.setValidating(true);
            dbf.setAttribute("http://xml.org/sax/features/validation",
                             Boolean.TRUE);
        }

        dbf.setNamespaceAware(true);
        dbf.setAttribute("http://xml.org/sax/features/namespaces", Boolean.TRUE);

        if (schemaValidate) {
            dbf.setAttribute("http://apache.org/xml/properties/schema/external-schemaLocation",
                             Constants.SignatureSpecNS + " " + signatureSchemaFile);
        }

        try {
            javax.xml.parsers.DocumentBuilder db = dbf.newDocumentBuilder();

            db.setErrorHandler(new org.apache.xml.security.utils.IgnoreAllErrorHandler());

            if (schemaValidate) {
                db.setEntityResolver(new org.xml.sax.EntityResolver() {

                    public org.xml.sax.InputSource resolveEntity(
                        String publicId, String systemId
                    ) throws org.xml.sax.SAXException {

                        if (systemId.endsWith("xmldsig-core-schema.xsd")) {
                            try {
                                return new org.xml.sax.InputSource(
                                    new FileInputStream(signatureSchemaFile));
                            } catch (FileNotFoundException ex) {
                                throw new org.xml.sax.SAXException(ex);
                            }
                        } else {
                            return null;
                        }
                    }
                });
            }
            // START COMPARING
            org.w3c.dom.Document doc = loadXMLFromString(xmlSignature, true);
            

            XPathFactory xpf = XPathFactory.newInstance();
            XPath xpath = xpf.newXPath();
            xpath.setNamespaceContext(new DSNamespaceContext());

            String expression = "//ds:Signature[1]";
            Element sigElement = 
                (Element) xpath.evaluate(expression, doc, XPathConstants.NODE);
            XMLSignature signature = 
                    new XMLSignature(sigElement, "");

            signature.addResourceResolver(new OfflineResolver());

            KeyInfo ki = signature.getKeyInfo();

            if (ki != null) {

                X509Certificate cert = signature.getKeyInfo().getX509Certificate();

                if (cert != null) {
                    return signature.checkSignatureValue(cert);
                } else {
                    PublicKey pk = signature.getKeyInfo().getPublicKey();

                    if (pk != null) {
                        return signature.checkSignatureValue(pk);
                    } else {
                        System.out.println(
                            "Did not find a public key, so I can't check the signature");
                    }
                }
            } else {
                System.out.println("Did not find a KeyInfo");
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
        return false;
    }
	
	/**
	 * Get the Certificate with the user's Identity from an XMLSignature
	 * @param xmlSignature
	 * @return
	 */
	public X509Certificate getCertificateFromXMLSignature(String xmlSignature) {
		X509Certificate certificate = null;
		try {
			Document doc = loadXMLFromString(xmlSignature, false);
			// XPathFactory erstellen
			XPathFactory xFactory = XPathFactory.newInstance();
			// XPath-Objekt erstellen
			XPath xpath = xFactory.newXPath();
			// XPath-Expression erzeugen
			XPathExpression expr = xpath.compile("/Signature/KeyInfo/X509Data/X509Certificate");
			String certString = (String) expr.evaluate(doc, XPathConstants.STRING);
			byte[] certBytes = Base64.decode(certString);
			ByteArrayInputStream bs = new ByteArrayInputStream(certBytes);
			CertificateFactory fac = CertificateFactory.getInstance("X.509");
			certificate = (X509Certificate) fac.generateCertificate(bs);
		} catch(Exception e) {
			e.printStackTrace();
		}
		
		return certificate;
	}
	
	/**
	 * Get the signed Object without its signature
	 * @param xmlSignature
	 * @return
	 */
	public String getSignedObjectFromXMlSignature(String xmlSignature) {
		String signedObject = "";
		try {
			Document doc = loadXMLFromString(xmlSignature, false);
			// XPathFactory erstellen
			XPathFactory xFactory = XPathFactory.newInstance();
			// XPath-Objekt erstellen
			XPath xpath = xFactory.newXPath();

			// XPath-Expression erzeugen
			XPathExpression expr = xpath.compile("/Signature");
			Node signatureNode = (Node) expr.evaluate(doc, XPathConstants.NODE);
			String signatureId = signatureNode.getAttributes().getNamedItem("Id").getTextContent().replaceFirst("signature", "");
			
			expr = xpath.compile("/Signature/Object[@Id=\"signed-data" + signatureId + "\"]");
			signedObject = (String) expr.evaluate(doc, XPathConstants.STRING);
			
		} catch(Exception e) {
			e.printStackTrace();
		}
		
		return signedObject;
	}
}
